#
# Copyright (c) 2011-2012 Johann Hanne
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#

import struct
import collections

LinkerFixup = collections.namedtuple('fixup', ('m', 'location',
    'data_record_offset', 'frame_method', 'frame_datum', 'target_method',
    'target_datum', 'target_displacement'))

class LinkerOutput(object):

    class chunk(object):

        class data(object):
            def __init__(self, data):
                self.data = bytearray(data)
                self.fixups = []

            def addfixup(self, **kwargs):
                self.fixups.append(LinkerFixup(**kwargs))

        def __init__(self, numbytes, objfilename, segmentname, segmentclassname, alignment, combination):
            if combination != 2:
                # We only support combination mode 2:
                # Combine by appending at an offset that meets the alignment requirement.
                raise RuntimeError("File {0}, segment {1}: Unsupported combination {2}".format(objfilename, segmentname, combination))

            self.numbytes = numbytes
            self.objfilename = objfilename
            self.segmentname = segmentname
            self.segmentclassname = segmentclassname
            self.alignment = alignment
            # Published names with associated offsets
            self.names = {}
            # Program data, per offset
            self.offsetdata = {}
            # Indicator if data is iterated (i.e. from LIDATA record), per offset
            self.is_iterated_data = {}
            self.outputpos = None
            self.group = None

        def __str__(self):
            return "<chunk object (segment name: '{0}')>".format(self.segmentname)

        def register_public_name(self, name, offset):
            self.names[name] = offset
            #print("Registered name '{0}' at {1}/{2}, offset {3}".format(name, self.objfilename, self.segmentname, offset))

        def setdata(self, offset, data, is_iterated_data):
            if offset in self.offsetdata:
                raise RuntimeError("{0}/{1}: Data for offset {2} already set".format(self.objfilename, self.segmentname, offset))
            self.offsetdata[offset] = self.data(data)
            self.is_iterated_data[offset] = is_iterated_data
            #print("Got {0} bytes of data for {1}/{2}, offset {3}".format(len(data), self.objfilename, self.segmentname, offset))
            return self.offsetdata[offset]

    class padchunk(object):

        class data(object):
            def __init__(self, padlen):
                self.data = bytearray('\xff' * padlen)
                self.fixups = []

        def __init__(self, padsize):
            self.offsetdata = { 0: self.data(padsize) }
            self.is_iterated_data = { 0: False }
            self.numbytes = padsize

    class chunkgroup(object):
        def __init__(self, objfilename, segmentgroupname):
            self.objfilename = objfilename
            self.segmentgroupname = segmentgroupname
            self.groupmembers = []

        def __str__(self):
            return "<chunkgroup object (segment name: '{0}')>".format(self.segmentgroupname)

        def is_member(self, x):
            return x in self.groupmembers

        def add_member(self, newmember):
            self.groupmembers.append(newmember)

    def __init__(self, filename, mapfilename, startsegment, startoffset, definitions):
        self.filename = filename
        self.mapfilename = mapfilename
        self.outpos = 0
        self.startaddress = (startsegment << 8) + startoffset
        self.definitions = definitions

        # List of segment names, in order of appearance
        self.segmentnames = []
        # List of allocated chunks per segment name
        self.chunks = {}

        # List of segment group names, in order of appearance
        self.segmentgroupnames = []
        # Segment group per segment group name
        self.segmentgroups = {}

        # Program data, per offset
        self.offsetdata = {}

        self.fixup_done = False

        self.got_bss = False

    def debugout(self, str):
        pass

    # Expands the iterated data from LIDATA records
    # (must be done *after* FIXUPs are applied)
    def get_expanded_iterated_data(self, data):
        offset = 0
        expanded_data = ""
        while offset < len(data):
            if (len(data) - offset) < 6:
                raise RuntimeError("Got LIDATA record with data block smaller than 6 bytes which is invalid")
            repeat_count = data[offset + 0] | (data[offset + 1] << 8)
            block_count = data[offset + 2] | (data[offset + 3] << 8)
            if block_count != 0:
                # We don't support recursive definitions as the OpenWatcom compiler doesn't seem to emit them anyway
                raise RuntimeError("Got LIDATA record data block with block count != 0 which is unsupported")
            count = data[offset + 4]
            content = data[offset + 5:offset + 5 + count]
            expanded_data += content * repeat_count
            offset += 5 + count
        return expanded_data

    def is_same_segment(self, segment1, segment2):
        if isinstance(segment1, LinkerOutput.chunk) and isinstance(segment2, LinkerOutput.chunk):
            return segment1 == segment2
        elif isinstance(segment1, LinkerOutput.chunkgroup) and isinstance(segment2, LinkerOutput.chunkgroup):
            return segment1 == segment2
        elif isinstance(segment1, LinkerOutput.chunkgroup):
            return segment1.is_member(segment2)
        else: # i.e. isinstance(segment2, LinkerOutput.chunkgroup):
            return segment2.is_member(segment1)

    def find_public_name(self, name):

        if name == '__bss_start':
            if self.got_bss:
                return (self.bss_segment, self.bss_offset)
            else:
                return (0, 0)
        elif name == '__bss_end':
            if self.got_bss:
                return (self.bss_segment, self.bss_offset + self.bss_size)
            else:
                return (0, 0)

        if name.startswith('__segment_'):
            segmentname = name[10:]
            if segmentname not in self.segmentnames:
                raise RuntimeError("Undefined reference to segment '{0}'".format(segmentname))
            chunk = self.chunks[segmentname][0]
            return (chunk.outputpos[0], chunk.outputpos[1])

        if name.startswith('__segmentgroup_'):
            segmentgroupname = name[15:]
            if segmentgroupname not in self.segmentgroupnames:
                raise RuntimeError("Undefined reference to segment group '{0}'".format(segmentgroupname))
            chunk = self.segmentgroups[segmentgroupname].groupmembers[0]
            return (chunk.outputpos[0], chunk.outputpos[1])

        for segmentname in self.segmentnames:
            for chunk in self.chunks[segmentname]:
                if name in chunk.names:
                    return (chunk.outputpos[0], chunk.outputpos[1] + chunk.names[name])

        if name in self.definitions:
            return self.definitions[name]

        raise RuntimeError("Undefined reference to '{0}'".format(name))

    def newsegment(self, objfilename, segmentname, segmentclassname, numbytes, alignment, combination):
        self.debugout("Allocating {0} bytes for {1}/{2}\n".format(numbytes, objfilename, segmentname))
        chunk = self.chunk(numbytes, objfilename, segmentname, segmentclassname, alignment, combination)
        if segmentname not in self.chunks:
            self.debugout("  First appearance of segment name {0}\n".format(segmentname))
            self.segmentnames.append(segmentname)
            self.chunks[segmentname] = []
        self.chunks[segmentname].append(chunk)
        return chunk

    def newsegmentgroup(self, objfilename, segmentgroupname, groupmembers):

        if segmentgroupname not in self.segmentgroups:
            self.debugout("Got segment group named '{0}' in {1} (first appearance); members are:\n".format(segmentgroupname, objfilename))
            chunkgroup = self.chunkgroup(objfilename, segmentgroupname)
            self.segmentgroupnames.append(segmentgroupname)
            self.segmentgroups[segmentgroupname] = chunkgroup
        else:
            self.debugout("Got segment group named '{0}' in {1}; members are:\n".format(segmentgroupname, objfilename))
            chunkgroup = self.segmentgroups[segmentgroupname]

        for groupmember in groupmembers:
            self.debugout("- {0}\n".format(groupmember.segmentname))
            if groupmember.group is not None:
                raise RuntimeError("Segment {0} is already a member of group {1}".format(groupmember.segmentname, groupmember.group.segmentgroupname))
            groupmember.group = chunkgroup
            chunkgroup.add_member(groupmember)

        return chunkgroup

    def write_le_s16(self, data, offset, val):
        if val < -32768 or val > 32767:
            raise RuntimeError("Got invalid value for signed 16-bit integer: {0}".format(val))
        newdata = struct.pack('<h', val)
        data[offset:offset + len(newdata)] = newdata

    def write_le_u16(self, data, offset, val):
        if val < 0 or val > 65535:
            raise RuntimeError("Got invalid value for unsigned 16-bit integer: {0}".format(val))
        newdata = struct.pack('<H', val)
        data[offset:offset + len(newdata)] = newdata

    def read_le_s16(self, data, offset):
        b = data[offset:offset + 2]
        return struct.unpack('<h', str(b))[0]

    def read_le_u16(self, data, offset):
        b = data[offset:offset + 2]
        return struct.unpack('<H', str(b))[0]

    def check_for_duplicate_names(self):
        name_to_chunk = {}

        for segmentname in self.segmentnames:
            for chunk in self.chunks[segmentname]:
                for name in chunk.names.iterkeys():
                    if name in name_to_chunk:
                        raise RuntimeError("Duplicate symbol: '{0}' ({1}, {2})".format(name, chunk.objfilename, name_to_chunk[name].objfilename))
                    name_to_chunk[name] = chunk

    def get_padsize(self, curaddress, alignment):
        if alignment == 2:
            # word (2-byte, 16-bit) aligned
            if (curaddress % 0x02) != 0:
                return ((curaddress | 0x01) + 1) - curaddress
        elif alignment == 3:
            # paragraph (256-byte) aligned
            # (on regular real mode, a paragraph has 16 bytes)
            if (curaddress % 0x100) != 0:
                return ((curaddress | 0xff) + 1) - curaddress
        else:
            raise RuntimeError("File {0}, segment {1}: Unsupported alignment {2}".format(chunk.objfilename, chunk.segmentname, chunk.alignment))

        return 0

    def fixup(self):
        self.debugout("Sanity check...\n")

        self.bss_size = 0

        for segmentname in self.segmentnames:
            self.debugout("- segment '{0}'\n".format(segmentname))

            for chunk in self.chunks[segmentname]:

                curchunksize = 0

                if chunk.group is None:
                    groupinfo = ""
                else:
                    groupinfo = " (puts this segment into group {0})".format(chunk.group.segmentgroupname)

                self.debugout("  - {0} bytes from {1}{2}\n".format(chunk.numbytes, chunk.objfilename, groupinfo))
                offsets = chunk.offsetdata.keys()
                offsets.sort()
                for offset in offsets:
                    self.debugout("    - offset {0}\n".format(offset))

                    if offset != curchunksize:
                        raise RuntimeError("Got non-contiguous data: expected offset {0}, got offset {1}".format(curchunksize, offset))

                    if chunk.is_iterated_data[offset]:
                        curchunksize += len(self.get_expanded_iterated_data(chunk.offsetdata[offset].data))
                    else:
                        curchunksize += len(chunk.offsetdata[offset].data)

                if chunk.segmentclassname == 'BSS':
                    if curchunksize != 0:
                        raise RuntimeError("Data length of BSS is not 0, but {0} bytes".format(curchunksize))
                else:
                    if curchunksize != chunk.numbytes:
                        raise RuntimeError("Data length inconsistency: object file indicated {0} bytes, but got {1} bytes".format(chunk.numbytes, curchunksize))

        self.debugout("Assigning addresses...\n")

        curaddress = self.startaddress

        for segmentname in self.segmentnames:

            self.debugout("- segment '{0}'\n".format(segmentname))

            for chunk in self.chunks[segmentname]:
                if chunk.group is not None:
                    self.debugout("  - skipping data from {0} as it is part of group {1} there\n".format(chunk.objfilename, chunk.group.segmentgroupname))
                    continue

                self.debugout("  - object file {0}\n".format(chunk.objfilename))

                if chunk.segmentclassname == 'BSS':
                    raise RuntimeError("BSS segment outside segment group not supported (yet?)")

                padsize = self.get_padsize(curaddress, chunk.alignment)
                if padsize != 0:
                    self.offsetdata[curaddress] = self.padchunk(padsize)
                    curaddress += padsize

                addr_segment = curaddress // 0x100
                addr_offset = curaddress - (addr_segment * 0x100)
                chunk.outputpos = (addr_segment, addr_offset)
                self.offsetdata[curaddress] = chunk
                self.debugout("    outputpos {0:x} ({1:x}:{2:x})\n".format(curaddress, addr_segment, addr_offset))
                curaddress += chunk.numbytes

        for segmentgroupname in self.segmentgroupnames:
            if self.got_bss:
                raise RuntimeError("BSS segment not part of *last* segment group - not (yet?) supported")

            groupsize = 0

            # All group members must share the same segment
            addr_segment = curaddress // 0x100
            self.segmentgroups[segmentgroupname].addr = addr_segment

            self.debugout("- segment group '{0}'\n".format(segmentgroupname))

            bss_chunks = []

            for chunk in self.segmentgroups[segmentgroupname].groupmembers:
                self.debugout("  - segment '{0}' of object file {1}\n".format(chunk.segmentname, chunk.objfilename))

                if chunk.segmentclassname == 'BSS':
                    self.debugout("    (BSS segment chunk delayed to group end)\n")
                    bss_chunks.append(chunk)
                    continue

                padsize = self.get_padsize(curaddress, chunk.alignment)
                if padsize != 0:
                    self.offsetdata[curaddress] = self.padchunk(padsize)
                    curaddress += padsize

                addr_offset = curaddress - (addr_segment * 0x100)
                chunk.outputpos = (addr_segment, addr_offset)
                self.offsetdata[curaddress] = chunk
                if chunk.numbytes != 0:
                    self.debugout("    - {0} bytes to {1:x} (0x{2:04x}:0x{3:04x})\n".format(chunk.numbytes, curaddress, addr_segment, addr_offset))
                curaddress += chunk.numbytes
                groupsize += chunk.numbytes

            if len(bss_chunks) > 0:
                self.bss_segment = addr_segment
                self.bss_offset = curaddress - (self.bss_segment * 0x100)
                self.got_bss = True

                self.debugout("  - collected BSS chunks:\n")

                for bss_chunk in bss_chunks:
                    padsize = self.get_padsize(curaddress, bss_chunk.alignment)
                    # No padchunk for BSS - the BSS is not part of the image,
                    # but memzeroed at boot!
                    curaddress += padsize
                    self.bss_size += padsize

                    addr_offset = curaddress - (addr_segment * 0x100)
                    bss_chunk.outputpos = (addr_segment, addr_offset)
                    self.offsetdata[curaddress] = bss_chunk
                    self.debugout("    - {0}: {1} bytes to {2:x} (0x{3:04x}:0x{4:04x})\n".format(bss_chunk.objfilename, bss_chunk.numbytes, curaddress, addr_segment, addr_offset))
                    curaddress += bss_chunk.numbytes
                    self.bss_size += bss_chunk.numbytes
                    groupsize += bss_chunk.numbytes

            if groupsize > (0x10000 - 0x100):
                # As the group members must be put into a single segment, the
                # combined size must not exceed 0x10000 (segment size) minus
                # 0x100 (paragraph size) bytes
                raise RuntimeError("Segment group too large")

        self.debugout("BSS size: {0} bytes\n".format(self.bss_size))

        self.debugout("Performing fixups...\n")

        for segmentname in self.segmentnames:
            self.debugout("- segment '{0}'\n".format(segmentname))

            for chunk in self.chunks[segmentname]:

                self.debugout("  {0} bytes\n".format(chunk.numbytes))
                offsets = chunk.offsetdata.keys()
                offsets.sort()
                for offset in offsets:
                    self.debugout("    offset {0:04x}\n".format(offset))

                    for fixup in chunk.offsetdata[offset].fixups:
                        self.debugout("       fixup target_method={0}\n".format(fixup.target_method))

                        # This will generate an error if these variables are
                        # accessed without being initialized before
                        addr_segment = None
                        addr_offset = None

                        if fixup.m == 0:
                            # M = 0 indicates a "self-relative fixup"

                            if fixup.target_method == 6:
                                # Specified by a EXTDEF index only; the displacement in the FIXUP subrecord is assumed to be 0.
                                (addr_segment, addr_offset) = self.find_public_name(fixup.target_datum)
                            else:
                                raise RuntimeError("Target method {0} not supported".format(fixup.target_method))

                            if fixup.location == 1:
                                # 16-bit offset
                                if fixup.frame_method == 5:
                                    # The FRAME is determined by the TARGET's segment, group, or external index.
                                    # (this seems to mean that we can rely on being in the target segment??)
                                    pass
                                else:
                                    raise RuntimeError("Location type 1 not supported for frame method {0}".format(fixup.frame_method))

                                if self.read_le_s16(chunk.offsetdata[offset].data, fixup.data_record_offset) != 0:
                                    raise RuntimeError("Got a self-relative fixup with a non-zero original offset?!")
                                    
                                if chunk.outputpos[0] != addr_segment:
                                    raise RuntimeError("Self-relative fixup refers to different segment")

                                # TODO: It's not entirely clear why the "- 2" is required - just because the
                                # offset number is always 16 bits large?
                                newdata = ((addr_offset) - (offset + chunk.outputpos[1] + fixup.data_record_offset)) - 2

                                self.write_le_s16(chunk.offsetdata[offset].data, fixup.data_record_offset, newdata)
                            else:
                                raise RuntimeError("Location type {0} not supported".format(fixup.location))

                        else:

                            if fixup.target_method == 4:
                                # Specified by a SEGDEF index only; the displacement in the FIXUP subrecord is assumed to be 0.
                                (addr_segment, addr_offset) = fixup.target_datum.outputpos
                            elif fixup.target_method == 6:
                                # Specified by a EXTDEF index only; the displacement in the FIXUP subrecord is assumed to be 0.
                                (addr_segment, addr_offset) = self.find_public_name(fixup.target_datum)
                            else:
                                raise RuntimeError("Target method {0} not supported".format(fixup.target_method))

                            if fixup.location == 1:
                                # 16-bit offset
                                if fixup.frame_method == 1:
                                    # The FRAME is specified by a GRPDEF index.
                                    if not self.is_same_segment(fixup.frame_datum, fixup.target_datum):
                                        raise RuntimeError("Fixup with 16-bit offset requested, but target not in same segment")
                                elif fixup.frame_method == 5:
                                    # The FRAME is determined by the TARGET's segment, group, or external index.
                                    # (this seems to mean that we can rely on being in the target segment??)
                                    pass
                                else:
                                    raise RuntimeError("Location type 1 not supported for frame method {0}".format(fixup.frame_method))
                                origdata = self.read_le_s16(chunk.offsetdata[offset].data, fixup.data_record_offset)
                                newdata = origdata + addr_offset
                                self.write_le_u16(chunk.offsetdata[offset].data, fixup.data_record_offset, newdata)
                            elif fixup.location == 2:
                                # 16-bit base
                                if fixup.frame_method == 1:
                                    # The FRAME is specified by a GRPDEF index.
                                    if not self.is_same_segment(fixup.frame_datum, fixup.target_datum):
                                        raise RuntimeError("Fixup with 16-bit offset requested, but target not in same segment")
                                elif fixup.frame_method == 5:
                                    # The FRAME is determined by the TARGET's segment, group, or external index.
                                    # (this seems to mean that we can rely on being in the target segment??)
                                    pass
                                else:
                                    raise RuntimeError("Location type 2 not supported for frame method {0}".format(fixup.frame_method))
                                origdata = self.read_le_s16(chunk.offsetdata[offset].data, fixup.data_record_offset)
                                if origdata != 0:
                                    raise RuntimeError("Got location type 2 with original data != 0, which doesn't seem to make sense?!")
                                newdata = addr_segment
                                self.write_le_u16(chunk.offsetdata[offset].data, fixup.data_record_offset, newdata)
                            elif fixup.location == 3:
                                # 32-bit Long pointer (16-bit base:16-bit offset).
                                # We don't care about frame_method here, as we are
                                # filling in an absolute address anyway
                                orig_offset = self.read_le_s16(chunk.offsetdata[offset].data, fixup.data_record_offset)
                                new_offset = orig_offset + addr_offset
                                self.write_le_u16(chunk.offsetdata[offset].data, fixup.data_record_offset, new_offset)
                                orig_segment = self.read_le_s16(chunk.offsetdata[offset].data, fixup.data_record_offset + 2)
                                if orig_segment != 0:
                                    raise RuntimeError("Got location type 3 with original segment != 0?!")
                                self.write_le_u16(chunk.offsetdata[offset].data, fixup.data_record_offset + 2, addr_segment)
                            else:
                                raise RuntimeError("Location type {0} not supported".format(fixup.location))

                    if 'target_displacement' in fixup and fixup.target_displacement != 0:
                        raise RuntimeError("target_displacement {0} field in fixup data not supported ({1})".format(fixup.target_displacement, chunk.objfilename))

                    curchunksize += len(chunk.offsetdata[offset].data)

        self.fixup_done = True

    def write_map(self):
        if self.mapfilename is None:
            return

        f = open(self.mapfilename, 'wt')

        offsets = self.offsetdata.keys()
        offsets.sort()
        for offset in offsets:
            chunk = self.offsetdata[offset]
            if isinstance(chunk, self.padchunk):
                f.write("  {0} bytes padding\n".format(chunk.numbytes))
                continue

            f.write("  {0} bytes from object file {1} to segment {2}.{3}\n".format(chunk.numbytes, chunk.objfilename, chunk.segmentname, chunk.segmentclassname))
            segment_offset_name = []
            for name in chunk.names:
                segment_offset_name.append((chunk.outputpos[0], chunk.outputpos[1] + chunk.names[name], name))
            for segment, offset, name in sorted(segment_offset_name):
                f.write("   {0:40s}0x{1:04x}:0x{2:04x}\n".format(name, segment, offset))

        f.close()

    def write(self):
        if not self.fixup_done:
            raise RuntimeError("write called before fixup")

        f = open(self.filename, 'wb')

        offsets = self.offsetdata.keys()
        offsets.sort()
        written = 0
        for offset in offsets:

            suboffsets = self.offsetdata[offset].offsetdata.keys()
            suboffsets.sort()
            for suboffset in suboffsets:

                if ((offset + suboffset) - offsets[0]) != written:
                    raise RuntimeError("Linker output inconsistency")

                if self.offsetdata[offset].is_iterated_data[suboffset]:
                    data = self.get_expanded_iterated_data(self.offsetdata[offset].offsetdata[suboffset].data)
                else:
                    data = self.offsetdata[offset].offsetdata[suboffset].data

                f.write(data)
                written += len(data)

        f.close()

